import os
import talib
import tushare as ts
# from matplotlib import pyplot as plt
import mplfinance as mpf
import pandas as pd  # 科学计算的一个库
from Final_assignment.tools import imageProcess, make_folder

pro = ts.pro_api('0e07eada964de0e2ca5ec7e52c85e5b5094dd75eb0da23792c45e39b')  # tushare接口
cycle_before = 20  # 重叠周期数


class Stock(object):  # 股票类
    def __init__(self, stock_name, cycle):  # 初始函数
        self.stock_name = stock_name  # 股票代码
        self.pro = pro  # ts.pro_api
        self.stock_data = pd.DataFrame()  # 股票数据，dataFrame 结构
        self.label_csv = pd.DataFrame(columns=['filename', 'date', 'label'])  # 标签csv
        self.picture_nums = 0
        self.data_folder = os.path.join("data", cycle + "_data\\")  # 股价文件夹
        self.Kline_folder = os.path.join("Kline_diagram", cycle + "_Kline", stock_name)  # k线图文件夹
        self.label_folder = os.path.join("Kline_diagram", cycle + "_label\\")  # 标记文件夹

        # 创建文件夹
        make_folder("data")  # 创建data文件夹，保存数据
        make_folder(self.data_folder)  # 创建周期(daily/weekly)价格文件夹
        make_folder("Kline_diagram")  # 创建Kline_diagram，保存图片和标签
        make_folder(os.path.join("Kline_diagram", cycle + "_Kline"))  # 创建周期(daily/weekly)图片文件夹
        make_folder(self.Kline_folder)  # 创建该股票k线图的文件夹
        make_folder(os.path.join("Kline_diagram", cycle + "_label/"))  # 创建周期(daily/weekly)标记文件夹

    # 获取日价行情，并保存数据
    def get_price_daily(self):
        try:
            old_stock_data = pd.read_csv(self.data_folder + self.stock_name + '.csv')
            self.label_csv = pd.read_csv(self.label_folder + self.stock_name + '.csv', index_col=0)
            old_end_date = str(old_stock_data.iloc[-1]["trade_date"])  # 取最后一行的日期
        except:  # FileNotFoundError
            old_end_date = ""
        self.stock_data = self.pro.daily(**{
            "ts_code": self.stock_name,
            "trade_date": "",
            "start_date": old_end_date,
            "end_date": "",
            "offset": "",
            "limit": ""
        }, fields=[
            "ts_code",  # 股票代码
            "trade_date",  # 交易日期
            "open",  # 开
            "high",  # 高
            "low",  # 低
            "close",  # 收
            "pre_close",  # 昨收
            "change",  # 涨跌额close - pre_close
            "pct_chg",  # 涨跌幅
            "vol",  # 成交量
            "amount"  # 成交额
        ])
        # self.stock_data.sort_values(by='trade_date',axis=0,ascending=True,inplace=True)  # 时间升序排列
        self.stock_data = self.stock_data[::-1]  # 逆序，即时间顺序
        self.stock_data.index = range(0, self.stock_data.shape[0])  # 重新设置索引
        if old_end_date != "":  # 有旧数据
            print(self.stock_name, "is loading old daily data...")
            self.stock_data = self.stock_data.iloc[1:]  # 有重复行,删除
            self.stock_data = pd.concat([old_stock_data.iloc[-1 * cycle_before:], self.stock_data], axis=0,
                                        ignore_index=True)  # 拼接两个dataFrame
        print(self.stock_name, "is loading new daily data...")
        self.stock_data.to_csv(self.data_folder + self.stock_name + '.csv')  # 写入数据
        return self.stock_data

    # 获取周价行情，并保存数据
    def get_price_weekly(self):
        try:
            old_stock_data = pd.read_csv(self.data_folder + self.stock_name + '.csv', index_col=0)
            self.label_csv = pd.read_csv(self.label_folder + self.stock_name + '.csv', index_col=0)
            old_end_date = str(old_stock_data.iloc[-1]["trade_date"])  # 取最后一行的日期
        except:  # FileNotFoundError
            old_end_date = ""
        self.stock_data = self.pro.weekly(**{
            "ts_code": self.stock_name,
            "trade_date": "",
            "start_date": old_end_date,
            "end_date": "",
            "offset": "",
            "limit": ""
        }, fields=[
            "ts_code",  # 股票代码
            "trade_date",  # 交易日期
            "open",  # 开
            "high",  # 高
            "low",  # 低
            "close",  # 收
            "pre_close",  # 昨收
            "change",  # 涨跌额close - pre_close
            "pct_chg",  # 涨跌幅
            "vol",  # 成交量
            "amount"  # 成交额
        ])
        # self.stock_data.sort_values(by='trade_date',axis=0,ascending=True,inplace=True)  # 时间升序排列
        self.stock_data = self.stock_data[::-1]  # 逆序，即时间顺序
        self.stock_data.index = range(0, self.stock_data.shape[0])  # 重新设置索引
        if old_end_date != "":  # 有旧数据
            print(self.stock_name, "is loading old weekly data...")
            self.stock_data = self.stock_data.iloc[1:]  # 有重复行,删除
            self.stock_data = pd.concat([old_stock_data.iloc[-1 * cycle_before:], self.stock_data], axis=0,
                                        ignore_index=True)  # 拼接两个dataFrame
        print(self.stock_name, "is loading new weekly data...")
        self.stock_data.to_csv(self.data_folder + self.stock_name + '.csv')  # 写入数据
        return self.stock_data

    # 根据数据识别形态
    def find_stock_form(self, form_list1, form_list2):
        # talib.CDLFUNC返回值 [-100, 0, 100]。0就是无模式，100就是识别了模式，-100就是反的识别
        print(self.stock_name, "is finding stock form...")
        for form_name in form_list1:  # 无penetration参数
            self.stock_data[form_name] = eval("talib.CDL" + form_name)(self.stock_data.open, self.stock_data.high,
                                                                       self.stock_data.low,
                                                                       self.stock_data.close, )  # 2CROWS
        for form_name in form_list2:
            self.stock_data[form_name] = eval("talib.CDL" + form_name)(self.stock_data.open, self.stock_data.high,
                                                                       self.stock_data.low, self.stock_data.close,
                                                                       penetration=0)  # Morning Star 晨星

    # 绘制并标记
    def draw_and_marker(self, form_list, kline_nums, cycles_forecast):
        print(self.stock_name, "is drawing Kline diagram...")
        for form_name in form_list:  # 遍历所有形态
            indexs = self.stock_data[self.stock_data[form_name] != 0].index  # 按条件取值，得到形态列表
        for index in range(0, self.stock_data.shape[0]):
            index_left = max(0, index - kline_nums + 1)
            data = self.stock_data[index_left:index + 1]  # 切片取数值，待画图的数据，
            # 根据形态后的股价高低给出是否要购买
            forecast_index = index + cycles_forecast
            for form_name in form_list:  # 遍历所有形态
                if self.stock_data.iloc[index][form_name] != 0:
                    if forecast_index < self.stock_data.shape[0]:  # 有标签
                        if self.stock_data.iloc[index][form_name] and self.stock_data.iloc[forecast_index]["close"] > \
                                self.stock_data.iloc[index]["close"]:       # 看涨且涨
                            buy_or_not = 1  # buy
                        else:
                            buy_or_not = 0  # not buy
                        # k线图的命名
                        figname = os.path.join(self.Kline_folder,
                                               self.stock_name + str(
                                                   self.stock_data.iloc[index]["trade_date"]) + ".png", )
                        if os.path.exists(figname) is False:  # 若图片不存在
                            self.picture_nums += 1
                            self.label_csv.loc[self.picture_nums] = [
                                self.stock_name + str(self.stock_data.iloc[index]["trade_date"]) + ".png",
                                self.stock_data.iloc[index]["trade_date"], buy_or_not]  # 保存标签信息
                    else:  # 没有标签
                        make_folder(os.path.join(self.Kline_folder + "_unlabel"))  # 创建无标签的图片文件夹
                        figname = os.path.join(self.Kline_folder + "_unlabel",
                                               self.stock_name + str(
                                                   self.stock_data.iloc[index]["trade_date"]) + ".png", )

                    self.draw_candle_chart(data, figname)  # 画k线图
                    imageProcess(figname)  # 处理k线图
                    break
        self.label_csv.to_csv(self.label_folder + self.stock_name + '.csv')  # 保存标签数据

    # 绘制k线图
    def draw_candle_chart(self, data, figname):  # 画一个k线图
        # 列属性重命名，必须大写开头
        data = data.rename(
            columns={'trade_date': 'Date', 'open': 'Open', 'close': 'Close', 'high': 'High', 'low': 'Low'})
        my_color = mpf.make_marketcolors(up="#ff0000", down="#00ff00", inherit=True)  # k线颜色
        my_style = mpf.make_mpf_style(marketcolors=my_color, edgecolor="#FFFFFF")  # k线图样式
        data.set_index('Date', inplace=True)  # 设置行索引，inplace表示原地修改，否则
        data.index = pd.DatetimeIndex(data.index)  # 行索引类型转化，行索引必须是pandas.DatetimeIndex
        # 绘图，data的数据类型必须是pandas.DataFrame
        mpf.plot(data, type='candle', show_nontrading=False, style=my_style, savefig=figname, figscale=0.5,
                 figratio=(20, 20))

# 测试集与训练集的划分
# https://www.jiqizhixin.com/articles/2019-01-23-2'''
